# main.py
from flow import create_proofmapper_flow, create_quick_interactive_flow
import argparse
import os
import asyncio
import time

async def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='ProofMapper - 自动生成安全评估证明材料清单（支持交互式修改）')
    parser.add_argument('pdf_path', nargs='?', help='安全评估指南PDF文件的路径（快速交互模式下可选）')
    parser.add_argument('--output', '-o', default='output.json', help='输出JSON文件的路径')
    parser.add_argument('--interactive', '-i', action='store_true', help='启用交互式修改模式')
    args = parser.parse_args()
    
    # 检测启动模式
    is_quick_mode = args.interactive and not args.pdf_path and os.path.exists(args.output)
    
    if args.interactive and not args.pdf_path and not os.path.exists(args.output):
        # 错误：想要快速交互模式但JSON文件不存在
        print(f"❌ 错误: 快速交互模式需要现有的JSON文件")
        print(f"📄 指定的JSON文件不存在: {args.output}")
        print("\n使用方法:")
        print("  完整处理: python main.py <pdf_path> [--interactive]")
        print("  快速交互: python main.py --interactive --output <existing_json>")
        return
    elif is_quick_mode:
        # 快速交互模式：直接从现有JSON开始
        print("🚀 启动ProofMapper快速交互模式...")
        await run_quick_interactive_mode(args)
    else:
        # 完整处理模式：从PDF开始
        if not args.pdf_path:
            print("❌ 错误: 必须提供PDF文件路径")
            print("\n使用方法:")
            print("  完整处理: python main.py <pdf_path> [--interactive]")
            print("  快速交互: python main.py --interactive --output <existing_json>")
            return
        
        if not os.path.exists(args.pdf_path):
            print(f"❌ 错误: PDF文件不存在: {args.pdf_path}")
            return
        
        await run_full_processing_mode(args)


async def run_full_processing_mode(args):
    """完整处理模式：从PDF开始的完整流程"""
    # 初始化共享存储 - 包含交互式修改相关数据
    shared = {
        "pdf_path": args.pdf_path,
        "pdf_text": "",
        "text_chunks": [],           # 文本块
        "chunk_checkpoints": {},     # 每个块的评估项
        "checkpoints": [],           # 合并后的评估项
        "materials": [],             # 生成的证明材料
        "final_output": {},          # 最终JSON输出
        "enable_interactive": args.interactive,  # 是否启用交互式修改
        "modification_log": [],      # 修改记录（最多3次）
        "output_path": args.output   # 输出文件路径
    }
    
    print("🚀 启动ProofMapper多线程并行处理...")
    print(f"📄 处理文件: {args.pdf_path}")
    if args.interactive:
        print("💬 交互式修改模式已启用")
    
    # 记录开始时间
    start_time = time.time()
    
    # 创建并运行异步流程
    proofmapper_flow = create_proofmapper_flow()
    await proofmapper_flow.run_async(shared)
    
    # 计算处理时间
    end_time = time.time()
    processing_time = end_time - start_time
    
    # 显示处理时间统计（文件保存已在节点中完成）
    print(f"\n📊 处理性能统计:")
    print(f"   - 总处理时间: {processing_time:.2f} 秒")
    print(f"   - 文本块数: {len(shared.get('text_chunks', []))}")
    print(f"   - 评估项数: {len(shared.get('checkpoints', []))}")
    
    # 如果是非交互模式，显示最终统计（final_output统一为字典对象）
    if not args.interactive:
        result_data = shared["final_output"]
        print(f"   - 证明材料清单数: {len(result_data.get('items', []))}")
        
        # 显示模块分布统计
        modules = {}
        for item in result_data.get('items', []):
            module = item.get('module', '未知模块')
            modules[module] = modules.get(module, 0) + 1
        
        for module, count in sorted(modules.items()):
            print(f"   - {module}: {count} 项")


async def run_quick_interactive_mode(args):
    """快速交互模式：直接从现有JSON开始"""
    # 初始化共享存储 - 仅包含必要数据
    shared = {
        "output_path": args.output,    # JSON文件路径
        "final_output": {},            # 加载的JSON数据
        "modification_log": [],        # 修改记录（最多3次）
        "enable_interactive": True     # 启用交互式修改
    }
    
    # 记录开始时间
    start_time = time.time()
    
    # 创建并运行快速交互流程
    quick_flow = create_quick_interactive_flow()
    await quick_flow.run_async(shared)
    
    # 计算处理时间
    end_time = time.time()
    processing_time = end_time - start_time
    
    print(f"\n📊 快速交互模式统计:")
    print(f"   - 总处理时间: {processing_time:.2f} 秒")

def run_main():
    """同步入口函数，用于启动异步main"""
    asyncio.run(main())

if __name__ == "__main__":
    run_main()